import torch
import torch.nn as nn
from torchvision.models import EfficientNet_V2_S_Weights, efficientnet_v2_s


class Model(nn.Module):

    def __init__(self):
        super(Model, self).__init__()

        pretrained = efficientnet_v2_s(weights=EfficientNet_V2_S_Weights.IMAGENET1K_V1)

        # 重新组装模型，只要特征抽取部分
        pretrained = torch.nn.Sequential(
            pretrained.features,
            pretrained.avgpool,
            nn.Flatten(start_dim=1)
        )

        # 锁定参数，不训练
        for param in pretrained.parameters():
            param.requires_grad_(False)

        pretrained.eval()
        self.pretrained = pretrained

        # 线性输出层，这部分要重新训练
        self.fc = nn.Sequential(
            nn.Linear(1280, 256),
            nn.ReLU(),
            nn.Linear(256, 256),
            nn.ReLU(),
            nn.Linear(256, 100)
        )

    def forward(self, x):
        with torch.no_grad():
            x = self.pretrained(x)

        return self.fc(x)


def getModel():
    return Model()

#
# if __name__ == '__main__':
#     print(getModel())